import os
import pickle
import sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(project_root)
# import matplotlib
# matplotlib.use('Agg')

import numpy as np
import matplotlib.pyplot as plt

from scipy.stats import binned_statistic_2d
from bisect import bisect_left
from dateutil import parser

from utils import *

NUM_PACKAGE = 630
X_RANGE = (-40, 40)
Y_RANGE = (0, 80)
DELTA_X = 0.2
DELTA_Y = 0.2

def find_closest_time(input_time_str, data_list):
    if len(data_list) == 0:
        return 0, 1000
    input_time = parser.isoparse(input_time_str)

    times = [item[0] for item in data_list]

    pos = bisect_left(times, input_time)

    if pos == 0:
        closest_time = times[0]
    elif pos == len(times):
        closest_time = times[-1]
    else:
        before = times[pos - 1]
        after = times[pos]
        closest_time = before if abs((before - input_time).total_seconds()) <= abs((after - input_time).total_seconds()) else after

    min_diff = abs((closest_time - input_time).total_seconds())

    min_index = times.index(closest_time)

    return min_index, min_diff

LiDAR_data = []
current_LiDAR_data_id = 1
def get_nearest_packages_points(input_time_str, LiDAR_path, num_package=NUM_PACKAGE):
    global LiDAR_data
    global current_LiDAR_data_id

    min_diff = 1000
    while min_diff > 3e-2:        
        min_index, min_diff = find_closest_time(input_time_str[0], LiDAR_data)
    
        half = num_package // 2
        start_index = max(0, min_index - half)
        end_index = min(min_index + half, len(LiDAR_data))

        if end_index - start_index < num_package:
            file_path = os.path.join(LiDAR_path, f"{current_LiDAR_data_id}.pkl")
            current_LiDAR_data_id += 1
            try:
                with open(file_path, 'rb') as f:
                    chunk = pickle.load(f)
                    del LiDAR_data[:len(LiDAR_data) // 2]
                    LiDAR_data.extend(chunk)
                    LiDAR_data.sort(key=lambda x: x[0])
            except Exception:
                raise TypeError("The target time is out of range")
            
            min_diff = 1000
            continue

        selected_packages = LiDAR_data[start_index:end_index]

    collected_points = []
    for _, package_data in selected_packages:
        collected_points.extend(package_data["points"])

    return collected_points

# LiDAR to ground: 0.68，LiDAR to radar: (2.488, 0.236, -1.548)
def align_point_cloud(data, x_offset=2.488, y_offset=0.236, z_offset=-1.548+0.68, x_range=Y_RANGE, y_range=X_RANGE):
    x = data[:, 0] + x_offset
    y = data[:, 1] + y_offset
    z = data[:, 2] + z_offset

    aligned_data = np.stack((x, y, z), axis=-1)

    mask = (x_range[0] <= x) & (x <= x_range[1]) & (y_range[0] <= y) & (y <= y_range[1])
    filtered_data = aligned_data[mask]

    return filtered_data

def filter_data(X, Y, Z, x_range, y_range, azi, ele, err):
    z_threshold = np.percentile(Z, 95)
    mask = (Z != 0) & \
        (Z <= 2.5) & \
        (X >= x_range[0]) & (X <= x_range[1]) & \
        (Y >= y_range[0]) & (Y <= y_range[1]) & ((azi > 0.8) & (ele > 0.85)) & (err < 0.1 * np.sqrt(X**2 + Y**2))
    return mask

def plot_height_difference(radar_data, aligned_points, output_path, delta_x=DELTA_X, delta_y=DELTA_Y, x_range=X_RANGE, y_range=Y_RANGE):
    azi_map = radar_data.get('aziLinearlityMap', None)
    ele_map = radar_data.get('eleLinearlityMap', None)
    error_map = radar_data.get('errorMap', None)
    target_x = radar_data.get('targetX', None)
    target_y = radar_data.get('targetY', None)
    target_z = radar_data.get('targetZ', None)
    dn = radar_data.get('DN', None)

    if not all([azi_map is not None, ele_map is not None, error_map is not None, target_x is not None, 
                target_y is not None, target_z is not None, dn is not None]):
        raise ".mat文件中缺少数据"
    
    X = target_x.flatten()
    Y = target_y.flatten()
    Z = target_z.flatten()
    DN = dn.flatten()
    azi = azi_map.flatten()
    ele = ele_map.flatten()
    err = error_map.flatten()

    mask = filter_data(X, Y, Z, x_range, y_range, azi, ele, err)
    X = X[mask]
    Y = Y[mask]
    Z = Z[mask]
    Z = Z - Z.min()
    DN = DN[mask]
    azi = azi[mask]
    ele = ele[mask]
    err = err[mask]

    x_bins = np.arange(x_range[0], x_range[1] + delta_x, delta_x)
    y_bins = np.arange(y_range[0], y_range[1] + delta_y, delta_y)

    lidar_x = -aligned_points[:, 1]
    lidar_y = aligned_points[:, 0]
    lidar_z = aligned_points[:, 2]
    lidar_z = lidar_z - lidar_z.min()

    mask = lidar_z < 2.5
    lidar_x = lidar_x[mask]
    lidar_y = lidar_y[mask]
    lidar_z = lidar_z[mask]

    lidar_max_z, _, _,_ = binned_statistic_2d(lidar_x, lidar_y, lidar_z, statistic='mean', bins=[x_bins, y_bins])
    
    radar_max_z, _, _,_ = binned_statistic_2d(X, Y, Z, statistic='mean', bins=[x_bins, y_bins])

    height_diff = lidar_max_z - radar_max_z

    lidar_max_z = np.nan_to_num(lidar_max_z, nan=0.0)
    radar_max_z = np.nan_to_num(radar_max_z, nan=0.0)
    height_diff = np.nan_to_num(height_diff, nan=0.0)

    height_diff[height_diff==0] = np.min(height_diff)

    # plot_histogram(height_diff)

    fig, axs = plt.subplots(1, 3, figsize=(18, 6))

    for ax in axs:
        ax.set_facecolor('black')

    cmap = 'jet'
    # LiDAR Max Height
    im1 = axs[0].imshow(lidar_max_z.T, extent=[x_range[0], x_range[1], y_range[0], y_range[1]],
                        origin='lower', aspect='auto', cmap=cmap)
    axs[0].set_title('LiDAR Max Height')
    axs[0].set_xlabel('X (meters)')
    axs[0].set_ylabel('Y (meters)')
    fig.colorbar(im1, ax=axs[0], fraction=0.046, pad=0.04, label='Height (m)')

    # Radar Max Height
    im2 = axs[1].imshow(radar_max_z.T, extent=[x_range[0], x_range[1], y_range[0], y_range[1]],
                        origin='lower', aspect='auto', cmap=cmap)
    axs[1].set_title('Radar Max Height')
    axs[1].set_xlabel('X (meters)')
    axs[1].set_ylabel('Y (meters)')
    fig.colorbar(im2, ax=axs[1], fraction=0.046, pad=0.04, label='Height (m)')

    # Height Difference
    im3 = axs[2].imshow(height_diff.T, extent=[x_range[0], x_range[1], y_range[0], y_range[1]],
                        origin='lower', aspect='auto', cmap=cmap, vmin=np.min(height_diff), vmax=np.max(height_diff))
    axs[2].set_title('Height Difference (LiDAR - Radar)')
    axs[2].set_xlabel('X (meters)')
    axs[2].set_ylabel('Y (meters)')
    fig.colorbar(im3, ax=axs[2], fraction=0.046, pad=0.04, label='Height Difference (m)')

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, facecolor='white', bbox_inches='tight', pad_inches=0)
    plt.close()
    # plt.show()

def plot_histogram(Z):
    Z = np.ravel(Z)
    Z = Z[Z != 0]
    plt.hist(Z, bins=50, color='blue', alpha=0.7, label='Z values')
    plt.axvline(Z.all(), color='red', linestyle='--', label='95th Percentile')
    plt.xlabel('Z values')
    plt.ylabel('Frequency')
    plt.title('Histogram of Z Values')
    plt.legend()
    plt.show()
    
if __name__ == "__main__":
    folder_LiDAR_path = "./workspace/data/LiDAR/cache"
    radar_folder_path = "./workspace/data/SAR"
    output_folder = "./workspace/image/SAR/diff"

    os.makedirs(output_folder, exist_ok=True)

    for root, _, files in os.walk(radar_folder_path):
        files.sort(key=natural_key)
        for file in files:
            if file.endswith(".mat"):
                mat_path = os.path.join(root, file)
                radar_data = read_mat_file(mat_path)

                target_time_str = radar_data.get('utcTimeChar', None)  # '2024-12-01T13:01:45.5620505Z'

                 # 拾取目标时间附近 num_package 个包，num_point = 125 * num_package

                collected_points = get_nearest_packages_points(target_time_str, folder_LiDAR_path)
                collected_points = np.array(collected_points)  # (num_package * 125, 3)

                aligned_points = align_point_cloud(collected_points)

                output_path = os.path.join(output_folder, file).replace("mat", "png")
                plot_height_difference(radar_data, aligned_points, output_path)

                # break
    